package com.zzyl.common.ai;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.utils.JsonUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.Arrays;

@Component
@Slf4j
public class AIModelInvoker {

    @Autowired
    private AIProperties aIProperties;

    // 定义一个方法aiInvoker，接收一个字符串prompt作为参数
    public String aiInvoker(String prompt) throws InputRequiredException {
//        String url = aIProperties.getUrl();
//        String model = aIProperties.getModel();
//        String authorization = aIProperties.getAuthorization();
//        Map messageMap = Map.of("role","user","content",prompt);
//        List messages = List.of(messageMap);
//        Map bodyMap = new HashMap();
//        bodyMap.put("model",model);
//        bodyMap.put("max_output_tokens",4000);
//        bodyMap.put("messages",messages);
//        bodyMap.put("stream",false);
//        String body= HttpUtil.createRequest(Method.POST, url)
//                    // 设置头
//                    .header("Authorization",authorization)
//                    // 设置body参数  query   form
//                    .body(JSONUtil.toJsonStr(bodyMap))
//                    .execute().body();
//
//            System.out.println(body);
//        // 把json字符串转成对象
//        JSONObject jsonObject = JSONUtil.parseObj(body);
//        String content = jsonObject.getByPath("choices[0].message.content", String.class);
////        String content = jsonObject.getByPath("message.content", String.class);
//        content = content.replace("```json", "");
//        content = content.replace("```", "");
//        return content;
        String content = null;
        try {
            GenerationResult result = callWithMessage(prompt);
            System.out.println(JsonUtils.toJson(result));
            content = result.getOutput().getChoices().get(0).getMessage().getContent();
            content = content.replace("```json", "");
            content = content.replace("```", "");
            return content;
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            // 使用日志框架记录异常信息
            System.err.println("An error occurred while calling the generation service: " + e.getMessage());
        }
        System.exit(0);
        return content;

    }

    public GenerationResult callWithMessage(String prompt) throws ApiException, NoApiKeyException, InputRequiredException {
        String url = aIProperties.getUrl();
        String model = aIProperties.getModel();
        String authorization = aIProperties.getAuthorization();
        Generation gen = new Generation();
        Message userMsg = Message.builder()
                .role(Role.USER.getValue())
                .content(prompt)
                .build();
        GenerationParam param = GenerationParam.builder()
                // 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey(authorization)
                // 此处以qwen-plus为例，可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
                .model(model)
                .messages(Arrays.asList(userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .build();
        return gen.call(param);
    }
}